import math,numpy
import sys,os,shutil
import heapq
import debugpy
import random
sys.path.append("/home/fd_chen/in_context_RL/spoc-robot-training")

from environment.spoc_objects import SPOCObject
from utils.debug_utils import enable_remote_debug
from utils.constants.stretch_initialization_utils import STRETCH_ENV_ARGS, AGENT_ROTATION_DEG, AGENT_MOVEMENT_CONSTANT
from environment.stretch_controller import StretchController
from matplotlib import pyplot as plt
from path_A_star_expert.actions import determine_turn_actions,determine_turn_actions_end,ACTIONS,CANDIDATE_R_ACTIONS,get_new_position,cost_distance
from path_A_star_expert.astar_utils import  save_navigation_frame ,save_data_json,filter_objects,target_types,get_position_rotation,get_top_down_map,load_procthor_houses
from path_A_star_expert.actions import calculate_angle,normalize_angle,search_best_action,calculate_distance
class Node:
    def __init__(self,index_x,index_z, cost, pind,agent_pose=None):
        self.index_x = index_x # index of node x
        self.index_z = index_z # index of node z
        self.cost = cost  # g cost of node
        self.pind = pind  # parent index of node
        self.agent_pose = agent_pose # agent poseß

class StarExpert: 
    def __init__(self,house_id=0,scene=None):
            
        self.id = house_id
        self.root_dir = f"./debug/house_{self.id}"
        if not os.path.exists(self.root_dir):
            os.makedirs(self.root_dir)

        STRETCH_ENV_ARGS['gpu_device'] = 0

        # scene 
        self.stretch_controller = StretchController(**STRETCH_ENV_ARGS)
        self.stretch_controller.reset(scene)
        # environment
        self.map_grid = STRETCH_ENV_ARGS['gridSize']
        self.reachable_positions = self.stretch_controller.get_reachable_positions()  
        self.bounds = self.stretch_controller.controller.last_event.metadata["sceneBounds"]["size"]# width and height of the scene
        self.motion = self.get_motion()
        self.obsmap_positions = self.obsmap()

        # task
        self.successuful_objects = []

        # agent_start
        self.start_node = self.agent_start()

        
    def reset_task(self, target: SPOCObject):

        self.target = target
        self.goal_node = self.target_goal()
        self.pathx, self.pathz = [], []
        self.shortest_actions = []
        self.shortest_frames = []
        self.shortest_positions = []
        self.early_stop = False
        self.build_task_dir()
    
    def save_task(self):
        self.successuful_objects.append(self.target['objectId'])
        save_data_json(self.shortest_actions,f"{self.metadir}/actions.json")
        save_data_json({**self.target._thor_obj}, f"{self.metadir}/target.json")
        save_data_json(self.successuful_objects,f"{self.root_dir}/successuful_objects.json")
    
    def build_task_dir(self):

        self.frame_dir = f"{self.root_dir}/{self.target['objectType']+'|'+self.target['objectId']}/frames"
        self.metadir = f"{self.root_dir}/{self.target['objectType']+'|'+self.target['objectId']}/meta"
        for dir in [self.frame_dir, self.metadir]:
            if os.path.exists(dir):
                shutil.rmtree(dir)
            os.makedirs(dir, exist_ok=True)

    def target_goal(self):
        self.target_params = {
            'x': self.stretch_controller.get_object_position(self.target['objectId'])['x'],
            'z': self.stretch_controller.get_object_position(self.target['objectId'])['z'],
        }
        index_x = round((self.target_params['x'] - self.min_x) / self.map_grid)
        index_z = round((self.target_params['z'] - self.min_z) / self.map_grid)
        object_pose = {"x":self.target_params['x'],"z":self.target_params['z'],'rotation':.0}
        return Node(index_x,index_z, 0.0, -1,object_pose)
    
    def agent_start(self):
        self.agent_start_params = {
            'x': self.stretch_controller.get_current_agent_full_pose()['position']['x'],
            'z': self.stretch_controller.get_current_agent_full_pose()['position']['z'],
        }
        index_x = round((self.agent_start_params['x'] - self.min_x) / self.map_grid)
        index_z = round((self.agent_start_params['z'] - self.min_z) / self.map_grid)
        agent_pose ={"x":self.agent_start_params['x'],"z":self.agent_start_params['z'],"rotation":self.stretch_controller.get_current_agent_full_pose()['rotation']['y']}
        return Node(index_x,index_z, 0.0, -1,agent_pose)


    def is_reachable(self, x, z):
        for rp in self.reachable_positions:
            distance = math.sqrt((rp["x"] - x) ** 2 + (rp["z"] - z) ** 2)
            if distance < 0.01:
                return True
        return False

    def obsmap(self, vis=True):  # obstacle map is the positions that are not indexed
        self.obsmap_positions = []
        self.min_x, self.max_x = min(rp["x"] for rp in self.reachable_positions), max(
            rp["x"] for rp in self.reachable_positions
        )
        self.min_z, self.max_z = min(rp["z"] for rp in self.reachable_positions), max(
            rp["z"] for rp in self.reachable_positions
        )
        self.index_xw, self.index_zw = round((self.max_x-self.min_x)/self.map_grid), round((self.max_z-self.min_z)/self.map_grid)
        self.min_x_index, self.max_x_index = 0, self.index_xw
        self.min_z_index, self.max_z_index = 0, self.index_zw

        for x_ind in range(round((self.max_x - self.min_x) / self.map_grid) + 1):
            for z_ind in range(round((self.max_z - self.min_z) / self.map_grid) + 1):
                x = self.min_x + x_ind * self.map_grid
                z = self.min_z + z_ind * self.map_grid
                if not self.is_reachable(x, z):
                    self.obsmap_positions.append({"x": x, "z": z})

        if vis:
            self.visualize_env()
        
        return self.obsmap_positions
    
    def get_motion(self):
        self.motion = [[-1, 0], [-1, 1], [0, 1], [1, 1],
                [1, 0], [1, -1], [0, -1], [-1, -1]]
        return self.motion

    def calc_index(self, node):
        return node.index_z * self.index_xw + node.index_x
    
    def u_cost(self,u):
        return math.hypot(u[0], u[1])
    
    def fvalue(self, node, n_goal):
        return node.cost + self.h(node, n_goal)
    
    def h(self, node:Node, n_goal:Node):
        return math.hypot(node.index_x - n_goal.index_x, node.index_z - n_goal.index_z)
    
    def extract_path(self, closed_set):
        if not self.early_stop:
            pathx, pathz = [round(self.goal_node.index_x * self.map_grid + self.min_x, 3)], [round(self.goal_node.index_z * self.map_grid + self.min_z, 3)]
            n_ind = self.calc_index(self.goal_node)
        else:
            pathx, pathz = [round(self.end_node.index_x * self.map_grid + self.min_x, 3)], [round(self.end_node.index_z * self.map_grid + self.min_z, 3)]
            n_ind = self.calc_index(self.end_node)
            
        while True:
            node = closed_set[n_ind]
            pathx.append(round(node.index_x*self.map_grid+self.min_x,3))
            pathz.append(round(node.index_z*self.map_grid+self.min_z,3))
            n_ind = node.pind

            if node == self.start_node:
                break

        pathx = list(reversed(pathx))
        pathz = list(reversed(pathz))
        return pathx, pathz


    def a_start_planner(self):# only get the shortest path 
        _cost_map = numpy.full((self.index_zw, self.index_xw), numpy.inf)
        self.open_set, self.closed_set = {}, {}
        q_priority = []
        self.open_set[self.calc_index(self.start_node)] = self.start_node
        heapq.heappush(q_priority,(self.fvalue(self.start_node, self.goal_node), self.calc_index(self.start_node)))

        while self.open_set:
            _, ind = heapq.heappop(q_priority)
            n_curr = self.open_set[ind]
            self.closed_set[ind] = n_curr
            self.open_set.pop(ind)
            # print("open_set", len(self.open_set))

            if self.check_distance(n_curr, self.goal_node, threhold=1.0):
                node1 = Node(round((self.shortest_path[0][0]-self.min_x)/self.map_grid), 
                             round((self.shortest_path[0][1]-self.min_z)/self.map_grid), 0.0, -1,{})
                
                node2 = Node(round((self.goal_node.agent_pose['x']-self.min_x)/self.map_grid), 
                              round((self.goal_node.agent_pose['z']-self.min_z)/self.map_grid), 0.0, -1,{})   

                if self.h(node1,self.goal_node) > self.h(node2,self.goal_node):
                    self.shortest_path.pop(0)
                plt.imshow(_cost_map, cmap="hot", interpolation="nearest")
                plt.savefig(f"{self.metadir}/cost_map.png")
                self.draw(name=f"{self.metadir}/path")
                return self.shortest_path

            for i in range(len(self.motion)): # one step 
                
                detla_x_index, detla_z_index = self.motion[i]
                node = Node(n_curr.index_x + detla_x_index, n_curr.index_z + detla_z_index,
                            n_curr.cost + self.u_cost(self.motion[i]), ind)# index parent node

                if not self.check_node(node):
                    continue

                n_ind = self.calc_index(node) # 拓展节点的index

                if n_ind not in self.closed_set:
                    if n_ind in self.open_set:
                        if self.open_set[n_ind].cost > node.cost:
                            self.open_set[n_ind].cost = node.cost
                            self.open_set[n_ind].pind = ind
                    else:
                        self.open_set[n_ind] = node
                        heapq.heappush(q_priority, (self.fvalue(node, self.goal_node), self.calc_index(node)))
                _cost_map[node.index_z, node.index_x] = node.cost

        plt.imshow(_cost_map, cmap="hot", interpolation="nearest")
        plt.savefig(f"{self.metadir}/cost_map.png")

    
    def check_distance(self, node1:Node, node2:Node, threhold=0.6):
        if math.hypot((node1.index_x - node2.index_x)*self.map_grid, (node1.index_z - node2.index_z)*self.map_grid) < threhold:
            self.early_stop = True
            self.end_node = node1
            self.pathx, self.pathz = self.extract_path(self.closed_set)
            self.shortest_path = [[x,z] for x, z in zip(self.pathx,self.pathz)]
            return True
        return False

    def check_node(self, node:Node, threhold=0.2):  # FIXME:add collision check and action check
        
        if node.index_x <= self.min_x_index or node.index_x >= self.max_x_index or \
                node.index_z <= self.min_z_index or node.index_z >= self.max_z_index:
            return False

        for obs in self.obsmap_positions:# fixme: add collision check get the best action
            distance = math.sqrt((obs["x"] - (node.index_x*self.map_grid+self.min_x)) ** 2 + (obs["z"] - (node.index_z*self.map_grid+self.min_z)) ** 2)
            if distance <= threhold:
                return False  # in collision
        return True
    
    def update_start_node(self, agent_pose:dict):
        node =Node(round((agent_pose['x'] - self.min_x) / self.map_grid), round((agent_pose['z'] - self.min_z) / self.map_grid), 0.0, -1,agent_pose)
        self.start_node = node


    def path_to_action(self,path):
        if len(path) == 0:
            raise ValueError("Path is empty")
        
        # if self.start_node.agent_pose['x'] == path[0][0] and self.start_node.agent_pose['z'] == path[0][1]: # the target should not be the same as the origin
        #     path.pop(0)
        
        action,next_agent_pose = search_best_action(self.start_node.agent_pose, path[0], path[1])
        self.shortest_actions.extend(action)
        self.update_start_node(next_agent_pose)
        return action

    def explore_path(self):
        while (calculate_distance(self.start_node.agent_pose, self.goal_node.agent_pose) > 1.0):
            print(calculate_distance(self.start_node.agent_pose, self.goal_node.agent_pose))
            path = self.a_start_planner()

            action = self.path_to_action(path)
            print(f"Action: {action}")
        self.save_task()

        
    def draw(self, name="./debug/init.png"):
        fig, ax = plt.subplots(1, 1)
                
        xs_obs = [obs["x"] for obs in self.obsmap_positions]
        zs_obs = [obs["z"] for obs in self.obsmap_positions]
        plt.plot(xs_obs, zs_obs, 'sk')
        plt.plot(self.pathx, self.pathz, '-r')
        plt.plot(self.start_node.index_x*self.map_grid+self.min_x, self.start_node.index_z*self.map_grid+self.min_z, 'sg')
        plt.plot(self.goal_node.index_x*self.map_grid+self.min_x, self.goal_node.index_z*self.map_grid+self.min_z, 'sb')
        
        plt.axis("equal")
        plt.show()
        plt.savefig(f"{name}.png")
        plt.close()

    def visualize_env(self, mode="unreachable"):
        fig, ax = plt.subplots(1, 1)
        
        xs = [rp["x"] for rp in self.reachable_positions]
        zs = [rp["z"] for rp in self.reachable_positions]
        ax.scatter(xs, zs, color="blue")
        
        if mode == "unreachable":
            xs_obs = [obs["x"] for obs in self.obsmap_positions]
            zs_obs = [obs["z"] for obs in self.obsmap_positions]
            ax.scatter(xs_obs, zs_obs, color="red")
            
        ax.set_xlabel("$x$")
        ax.set_ylabel("$z$")
        ax.set_title("Positions in the Scene")
        ax.set_aspect("equal")
        plt.savefig(f"{self.root_dir}/positions_in_the_Scence.png")
        plt.close()

if __name__ == '__main__':
    enable_remote_debug(65532)
    houses = load_procthor_houses()
    house_id = 2
    # for house in houses:
    star_expert = StarExpert(house_id,scene=houses[house_id])
    objects = filter_objects(star_expert.stretch_controller,target_types())
    for target in objects:
        print(f"Target object: {target['objectId']},{target['objectType']}")
        star_expert.reset_task(target)
        star_expert.explore_path()
        star_expert.save_task()

    star_expert.stretch_controller.stop()
    house_id += 1
    